use axum::response::sse::{Event, KeepAlive};
use axum::{
    extract::{Path, State},
    http::StatusCode,
    response::{Json, Sse},
};
use futures::stream::Stream;
use rig_app::KnowledgeBaseConfig;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::convert::Infallible;
use std::sync::Arc;
use std::time::Duration;
use tokio::task;
use tokio_stream::StreamExt as _;

// Request/Response structs for agent management
#[derive(Deserialize)]
pub struct CreateAgentRequest {
    pub name: String,
    pub provider: String,
    pub model: String,
    pub preamble: String,
    pub base_url: Option<String>,
    pub tools: Option<Vec<String>>,
    pub api_key: Option<String>,
    pub category: Option<String>, // "chat", "embed", "image"
    pub knowledge_base: Option<KnowledgeBaseConfig>,
}

#[derive(Serialize)]
pub struct CreateAgentResponse {
    pub id: i64,
    pub message: String,
}

#[derive(Deserialize)]
pub struct ChatRequest {
    pub agent_id: i64,
    pub prompt: String,
    pub history: Vec<rig_app::ChatMessage>,
}

#[derive(Serialize)]
pub struct ChatResponse {
    pub message: String,
}

#[derive(Deserialize)]
pub struct PromptRequest {
    pub agent_id: i64,
    pub prompt: String,
}

#[derive(Serialize)]
pub struct PromptResponse {
    pub message: String,
}

#[derive(Serialize)]
pub struct AgentListResponse {
    pub agents: Vec<rig_app::AgentConfig>,
}

#[derive(Serialize)]
pub struct DeleteAgentResponse {
    pub success: bool,
    pub message: String,
}

#[derive(Deserialize)]
pub struct StreamChatRequest {
    pub agent_id: i64,
    pub prompt: String,
    pub history: Vec<rig_app::ChatMessage>,
}

#[derive(Serialize)]
pub struct StreamChunk {
    pub content: String,
    pub done: bool,
}

// MCP Service management request/response structs
#[derive(Deserialize)]
pub struct CreateMCPServiceRequest {
    pub name: String,
    pub protocol: String,
    pub command: Option<String>,
    pub args: Option<Vec<String>>,
    pub url: Option<String>,
    pub headers: Option<HashMap<String, String>>,
    pub env_vars: Option<HashMap<String, String>>,
    pub enabled: bool,
}

#[derive(Serialize)]
pub struct CreateMCPServiceResponse {
    pub id: i64,
    pub message: String,
}

#[derive(Serialize)]
pub struct MCPServiceListResponse {
    pub services: Vec<rig_app::MCPServiceConfig>,
}

#[derive(Serialize)]
pub struct MCPToolListResponse {
    pub tools: Vec<rig_app::MCPToolDefinition>,
}

#[derive(Serialize)]
pub struct DeleteMCPServiceResponse {
    pub success: bool,
    pub message: String,
}

#[derive(Deserialize)]
pub struct ExecuteMCPToolRequest {
    pub service_id: i64,
    pub tool_name: String,
    pub arguments: serde_json::Value,
}

#[derive(Serialize)]
pub struct ExecuteMCPToolResponse {
    pub success: bool,
    pub result: Option<serde_json::Value>,
    pub error: Option<String>,
}

// Agent management routes
pub async fn create_agent_handler(
    State(state): State<Arc<crate::AppState>>,
    Json(payload): Json<CreateAgentRequest>,
) -> Result<Json<CreateAgentResponse>, (StatusCode, String)> {
    let agent_manager = state.agent_manager.clone();

    task::spawn_blocking(move || {
        // Parse category from string
        let category = match payload.category.as_deref() {
            Some("embed") => rig_app::AgentCategory::Embed,
            Some("image") => rig_app::AgentCategory::Image,
            _ => rig_app::AgentCategory::Chat, // Default to Chat
        };

        match agent_manager.create_agent_sync(
            &payload.name,
            &payload.provider,
            &payload.model,
            &payload.preamble,
            payload.base_url.as_deref(),
            payload.tools.as_deref(),
            category,
            payload.api_key.as_deref(),
        ) {
            Ok(id) => Ok(Json(CreateAgentResponse {
                id,
                message: format!("Agent '{}' created successfully", payload.name),
            })),
            Err(e) => Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Failed to create agent: {}", e),
            )),
        }
    })
    .await
    .unwrap()
}

pub async fn list_agents_handler(
    State(state): State<Arc<crate::AppState>>,
) -> Result<Json<AgentListResponse>, (StatusCode, String)> {
    let agent_manager = state.agent_manager.clone();

    task::spawn_blocking(move || match agent_manager.list_agents_sync() {
        Ok(agents) => {
            // Convert EnhancedAgentConfig to AgentConfig for backward compatibility
            let agent_configs: Vec<rig_app::AgentConfig> = agents
                .into_iter()
                .map(|enhanced_agent| enhanced_agent.base_config)
                .collect();
            Ok(Json(AgentListResponse {
                agents: agent_configs,
            }))
        }
        Err(e) => Err((
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Failed to list agents: {}", e),
        )),
    })
    .await
    .unwrap()
}

pub async fn get_agent_handler(
    State(state): State<Arc<crate::AppState>>,
    Path(agent_id): Path<i64>,
) -> Result<Json<rig_app::AgentConfig>, (StatusCode, String)> {
    let agent_manager = state.agent_manager.clone();

    task::spawn_blocking(move || match agent_manager.get_agent_sync(agent_id) {
        Ok(Some(agent)) => Ok(Json(agent.base_config)),
        Ok(None) => Err((
            StatusCode::NOT_FOUND,
            format!("Agent with id {} not found", agent_id),
        )),
        Err(e) => Err((
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Failed to get agent: {}", e),
        )),
    })
    .await
    .unwrap()
}

pub async fn delete_agent_handler(
    State(state): State<Arc<crate::AppState>>,
    Path(agent_id): Path<i64>,
) -> Result<Json<DeleteAgentResponse>, (StatusCode, String)> {
    let agent_manager = state.agent_manager.clone();

    task::spawn_blocking(move || match agent_manager.delete_agent_sync(agent_id) {
        Ok(true) => Ok(Json(DeleteAgentResponse {
            success: true,
            message: format!("Agent with id {} deleted successfully", agent_id),
        })),
        Ok(false) => Err((
            StatusCode::NOT_FOUND,
            format!("Agent with id {} not found", agent_id),
        )),
        Err(e) => Err((
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Failed to delete agent: {}", e),
        )),
    })
    .await
    .unwrap()
}

// Chat functionality routes
pub async fn prompt_handler(
    State(state): State<Arc<crate::AppState>>,
    Json(payload): Json<PromptRequest>,
) -> Result<Json<PromptResponse>, (StatusCode, String)> {
    let agent_manager = state.agent_manager.clone();
    let tool_manager = state.tool_manager.clone();
    let prompt_request = payload;

    task::spawn_blocking(move || {
        tokio::runtime::Runtime::new().unwrap().block_on(async {
            match rig_app::prompt(
                &agent_manager,
                &tool_manager,
                prompt_request.agent_id,
                &prompt_request.prompt,
            )
            .await
            {
                Ok(response) => Ok(Json(PromptResponse { message: response })),
                Err(e) => Err((
                    StatusCode::INTERNAL_SERVER_ERROR,
                    format!("Failed to execute prompt: {}", e),
                )),
            }
        })
    })
    .await
    .unwrap()
}

pub async fn chat_handler(
    State(state): State<Arc<crate::AppState>>,
    Json(payload): Json<ChatRequest>,
) -> Result<Json<ChatResponse>, (StatusCode, String)> {
    let agent_manager = state.agent_manager.clone();
    let tool_manager = state.tool_manager.clone();
    let chat_request = payload;

    task::spawn_blocking(move || {
        tokio::runtime::Runtime::new().unwrap().block_on(async {
            match rig_app::chat(
                &agent_manager,
                &tool_manager,
                chat_request.agent_id,
                &chat_request.prompt,
                chat_request.history,
            )
            .await
            {
                Ok(response) => Ok(Json(ChatResponse { message: response })),
                Err(e) => Err((
                    StatusCode::INTERNAL_SERVER_ERROR,
                    format!("Failed to execute chat: {}", e),
                )),
            }
        })
    })
    .await
    .unwrap()
}

pub async fn stream_chat_handler(
    State(state): State<Arc<crate::AppState>>,
    Json(payload): Json<StreamChatRequest>,
) -> Result<Sse<impl Stream<Item = Result<Event, Infallible>>>, (StatusCode, String)> {
    let agent_manager = state.agent_manager.clone();
    let tool_manager = state.tool_manager.clone();
    let prompt = payload.prompt.clone();
    let history = payload.history.clone();
    let agent_id = payload.agent_id;

    let stream_result = task::spawn_blocking(move || {
        tokio::runtime::Runtime::new().unwrap().block_on(async {
            rig_app::stream_chat(&agent_manager, &tool_manager, agent_id, &prompt, history).await
        })
    })
    .await
    .unwrap();

    match stream_result {
        Ok(mut stream) => {
            let response_stream = async_stream::stream! {
                // Send initial event to establish connection
                yield Ok(Event::default().data("{\"type\": \"start\", \"message\": \"Stream started\"}"));

                while let Some(chunk_result) = stream.next().await {
                    match chunk_result {
                        Ok(content) => {
                            let chunk_data = StreamChunk {
                                content,
                                done: false,
                            };

                            match serde_json::to_string(&chunk_data) {
                                Ok(json_data) => {
                                    yield Ok(Event::default().data(json_data));
                                }
                                Err(e) => {
                                    let error_chunk = StreamChunk {
                                        content: format!("Serialization error: {}", e),
                                        done: true,
                                    };
                                    if let Ok(error_json) = serde_json::to_string(&error_chunk) {
                                        yield Ok(Event::default().data(error_json));
                                    }
                                    break;
                                }
                            }
                        }
                        Err(e) => {
                            let error_chunk = StreamChunk {
                                content: format!("Stream error: {}", e),
                                done: true,
                            };
                            if let Ok(error_json) = serde_json::to_string(&error_chunk) {
                                yield Ok(Event::default().data(error_json));
                            }
                            break;
                        }
                    }
                }

                // Send final "done" message
                let done_chunk = StreamChunk {
                    content: String::new(),
                    done: true,
                };
                if let Ok(done_json) = serde_json::to_string(&done_chunk) {
                    yield Ok(Event::default().data(done_json));
                }

                // Send final event to indicate stream completion
                yield Ok(Event::default().data("{\"type\": \"end\", \"message\": \"Stream completed\"}"));
            };

            Ok(Sse::new(response_stream).keep_alive(
                KeepAlive::new()
                    .interval(Duration::from_secs(15))
                    .text("keep-alive"),
            ))
        }
        Err(e) => Err((
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Failed to start stream chat: {}", e),
        )),
    }
}

// MCP Service management routes
#[axum::debug_handler]
pub async fn create_mcp_service_handler(
    State(state): State<Arc<crate::AppState>>,
    Json(payload): Json<CreateMCPServiceRequest>,
) -> Result<Json<CreateMCPServiceResponse>, (StatusCode, String)> {
    let tool_manager = state.tool_manager.clone();
    let payload = payload;

    task::spawn_blocking(move || {
        tokio::runtime::Runtime::new().unwrap().block_on(async {
            // Parse protocol string to enum
            let protocol = match payload.protocol.as_str() {
                "stdio" => rig_app::MCPProtocol::Stdio,
                "sse" => rig_app::MCPProtocol::Sse,
                "http" => rig_app::MCPProtocol::Http,
                _ => {
                    return Err((
                        StatusCode::BAD_REQUEST,
                        format!("Invalid protocol: {}", payload.protocol),
                    ))
                }
            };

            let service_config = rig_app::MCPServiceConfig {
                id: 0, // Will be assigned by register_mcp_service
                name: payload.name,
                protocol,
                command: payload.command,
                args: payload.args,
                url: payload.url,
                headers: payload.headers,
                env_vars: payload.env_vars,
                enabled: payload.enabled,
                created_at: chrono::Utc::now().to_rfc3339(),
            };

            match tool_manager.register_mcp_service(&service_config).await {
                Ok(id) => Ok(Json(CreateMCPServiceResponse {
                    id,
                    message: format!("MCP service '{}' created successfully", service_config.name),
                })),
                Err(e) => Err((
                    StatusCode::INTERNAL_SERVER_ERROR,
                    format!("Failed to create MCP service: {}", e),
                )),
            }
        })
    })
    .await
    .unwrap()
}

pub async fn list_mcp_services_handler(
    State(state): State<Arc<crate::AppState>>,
) -> Result<Json<MCPServiceListResponse>, (StatusCode, String)> {
    let tool_manager = state.tool_manager.clone();

    task::spawn_blocking(move || match tool_manager.list_mcp_services() {
        Ok(services) => Ok(Json(MCPServiceListResponse { services })),
        Err(e) => Err((
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Failed to list MCP services: {}", e),
        )),
    })
    .await
    .unwrap()
}

pub async fn get_mcp_service_handler(
    State(state): State<Arc<crate::AppState>>,
    Path(service_id): Path<i64>,
) -> Result<Json<rig_app::MCPServiceConfig>, (StatusCode, String)> {
    let tool_manager = state.tool_manager.clone();

    task::spawn_blocking(move || match tool_manager.get_mcp_service(service_id) {
        Ok(Some(service)) => Ok(Json(service)),
        Ok(None) => Err((
            StatusCode::NOT_FOUND,
            format!("MCP service with id {} not found", service_id),
        )),
        Err(e) => Err((
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Failed to get MCP service: {}", e),
        )),
    })
    .await
    .unwrap()
}

pub async fn delete_mcp_service_handler(
    State(state): State<Arc<crate::AppState>>,
    Path(service_id): Path<i64>,
) -> Result<Json<DeleteMCPServiceResponse>, (StatusCode, String)> {
    let tool_manager = state.tool_manager.clone();

    task::spawn_blocking(move || match tool_manager.delete_mcp_service(service_id) {
        Ok(true) => Ok(Json(DeleteMCPServiceResponse {
            success: true,
            message: format!("MCP service with id {} deleted successfully", service_id),
        })),
        Ok(false) => Err((
            StatusCode::NOT_FOUND,
            format!("MCP service with id {} not found", service_id),
        )),
        Err(e) => Err((
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Failed to delete MCP service: {}", e),
        )),
    })
    .await
    .unwrap()
}

pub async fn list_mcp_tools_handler(
    State(state): State<Arc<crate::AppState>>,
) -> Result<Json<MCPToolListResponse>, (StatusCode, String)> {
    let tool_manager = state.tool_manager.clone();

    task::spawn_blocking(move || match tool_manager.list_mcp_tools(None) {
        Ok(tools) => Ok(Json(MCPToolListResponse { tools })),
        Err(e) => Err((
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Failed to list MCP tools: {}", e),
        )),
    })
    .await
    .unwrap()
}

pub async fn execute_mcp_tool_handler(
    State(state): State<Arc<crate::AppState>>,
    Json(payload): Json<ExecuteMCPToolRequest>,
) -> Result<Json<ExecuteMCPToolResponse>, (StatusCode, String)> {
    let tool_manager = state.tool_manager.clone();
    let payload = payload;

    task::spawn_blocking(move || {
        tokio::runtime::Runtime::new().unwrap().block_on(async {
            match tool_manager
                .execute_mcp_tool(payload.service_id, &payload.tool_name, payload.arguments)
                .await
            {
                Ok(response) => Ok(Json(ExecuteMCPToolResponse {
                    success: response.success,
                    result: response.result,
                    error: response.error,
                })),
                Err(e) => Err((
                    StatusCode::INTERNAL_SERVER_ERROR,
                    format!("Failed to execute MCP tool: {}", e),
                )),
            }
        })
    })
    .await
    .unwrap()
}

// Enhanced functionality endpoints

#[derive(Deserialize)]
pub struct EmbeddingRequest {
    pub agent_id: i64,
    pub text: String,
}

#[derive(Serialize)]
pub struct EmbeddingResponse {
    pub embedding: Vec<f32>,
    pub dimensions: usize,
}

#[derive(Deserialize)]
pub struct ImageGenerationRequest {
    pub agent_id: i64,
    pub prompt: String,
    pub num_images: Option<u32>,
    pub width: Option<u32>,
    pub height: Option<u32>,
    pub quality: Option<String>,
}

#[derive(Serialize)]
pub struct ImageGenerationResponse {
    pub images: Vec<rig_app::GeneratedImage>,
}

#[derive(Deserialize)]
pub struct ImageAnalysisRequest {
    pub agent_id: i64,
    pub image_data: String,
    pub image_format: String,
    pub prompt: Option<String>,
    pub additional_params: Option<HashMap<String, serde_json::Value>>,
}

#[derive(Serialize)]
pub struct ImageAnalysisResponse {
    pub result: String,
    pub confidence: Option<f32>,
    pub metadata: HashMap<String, serde_json::Value>,
}

#[derive(Deserialize)]
pub struct RAGQueryRequest {
    pub agent_id: i64,
    pub query: String,
    pub limit: Option<usize>,
}

#[derive(Serialize)]
pub struct RAGQueryResponse {
    pub results: Vec<rig_app::EmbeddingSearchResult>,
    pub response: String,
}

// Embedding endpoint (simplified - requires embedding manager in AppState)
pub async fn create_embedding_handler(
    State(_state): State<Arc<crate::AppState>>,
    Json(_payload): Json<EmbeddingRequest>,
) -> Result<Json<EmbeddingResponse>, (StatusCode, String)> {
    Err((
        StatusCode::NOT_IMPLEMENTED,
        "Embedding creation requires EmbeddingManager in AppState".to_string(),
    ))
}

// Image generation endpoint (simplified - requires image manager in AppState)
pub async fn generate_image_handler(
    State(_state): State<Arc<crate::AppState>>,
    Json(_payload): Json<ImageGenerationRequest>,
) -> Result<Json<ImageGenerationResponse>, (StatusCode, String)> {
    Err((
        StatusCode::NOT_IMPLEMENTED,
        "Image generation requires ImageManager in AppState".to_string(),
    ))
}

// Image analysis endpoint (simplified - requires image manager in AppState)
pub async fn analyze_image_handler(
    State(_state): State<Arc<crate::AppState>>,
    Json(_payload): Json<ImageAnalysisRequest>,
) -> Result<Json<ImageAnalysisResponse>, (StatusCode, String)> {
    Err((
        StatusCode::NOT_IMPLEMENTED,
        "Image analysis requires ImageManager in AppState".to_string(),
    ))
}

// RAG query endpoint (simplified - requires embedding manager in AppState)
pub async fn rag_query_handler(
    State(_state): State<Arc<crate::AppState>>,
    Json(_payload): Json<RAGQueryRequest>,
) -> Result<Json<RAGQueryResponse>, (StatusCode, String)> {
    Err((
        StatusCode::NOT_IMPLEMENTED,
        "RAG query requires EmbeddingManager in AppState".to_string(),
    ))
}
